package podman

import (
	"github.com/containers/podman/v5/pkg/bindings/containers"
	"github.com/containers/podman/v5/pkg/specgen"
)

const (
	NvidiaGPUParent = "nvidia.com/gpu.parent"
	NvidiaGpuCount  = "nvidia.com/gpu.count"
)

var (
	EnvNvidiaVisibleDevicesAll = Env{
		Key: "NVIDIA_VISIBLE_DEVICES",
		Val: "all",
	}
	EnvNvidiaDriverCapabilities = Env{
		Key: "NVIDIA_DRIVER_CAPABILITIES",
		Val: "compute,utility",
	}
)

type Label struct {
	Tag, Val string
}

type Env struct {
	Key, Val string
}

type Annotation struct {
	Key, Val string
}

type SpecBuilder struct {
	spec *specgen.SpecGenerator
}

func NewSpec(rowImage string) *SpecBuilder {
	return &SpecBuilder{
		spec: specgen.NewSpecGenerator(rowImage, false),
	}
}

func (sb *SpecBuilder) WithLabel(label Label) *SpecBuilder {
	if sb.spec.Labels == nil {
		sb.spec.Labels = make(map[string]string)
	}
	sb.spec.Labels[label.Tag] = label.Val
	return sb
}
func (sb *SpecBuilder) WithLabels(labels []Label) *SpecBuilder {
	for _, label := range labels {
		sb.WithLabel(label)
	}
	return sb
}

func (sb *SpecBuilder) WithEnv(env Env) *SpecBuilder {
	if sb.spec.Env == nil {
		sb.spec.Env = make(map[string]string)
	}
	sb.spec.Env[env.Key] = env.Val
	return sb
}
func (sb *SpecBuilder) WithEnvs(envs []Env) *SpecBuilder {
	for _, env := range envs {
		sb.WithEnv(env)
	}
	return sb
}

func (sb *SpecBuilder) WithAnnotation(annotation Annotation) *SpecBuilder {
	if sb.spec.Annotations == nil {
		sb.spec.Annotations = make(map[string]string)
	}
	sb.spec.Annotations[annotation.Key] = annotation.Val
	return sb
}
func (sb *SpecBuilder) WithAnnotations(annotations []Annotation) *SpecBuilder {
	for _, annotation := range annotations {
		sb.WithAnnotation(annotation)
	}
	return sb
}

func (sb *SpecBuilder) WithName(name string) *SpecBuilder {
	sb.spec.Name = name
	return sb
}

func (sb *SpecBuilder) WithCommand(command ...string) *SpecBuilder {
	sb.spec.Command = command
	return sb
}

func (sb *SpecBuilder) Build() *specgen.SpecGenerator {
	return sb.spec
}

type OptBuilder struct {
	startOpts  *containers.StartOptions
	listOpts   *containers.ListOptions
	removeOpts *containers.RemoveOptions
	stopOpts   *containers.StopOptions
}

func NewStartOpts() *OptBuilder {
	return &OptBuilder{
		startOpts: new(containers.StartOptions),
	}
}
func NewListOpts() *OptBuilder {
	return &OptBuilder{
		listOpts: new(containers.ListOptions),
	}
}

func NewRemoveOpts() *OptBuilder {
	return &OptBuilder{
		removeOpts: new(containers.RemoveOptions),
	}
}

func NewStopOpts() *OptBuilder {
	return &OptBuilder{
		stopOpts: new(containers.StopOptions),
	}
}

func (ob *OptBuilder) BuildStartOptions() *containers.StartOptions {
	return ob.startOpts
}
func (ob *OptBuilder) BuildListOptions() *containers.ListOptions {
	return ob.listOpts
}

func (ob *OptBuilder) BuildRemoveOptions() *containers.RemoveOptions {
	return ob.removeOpts
}

func (ob *OptBuilder) BuildStopOptions() *containers.StopOptions {
	return ob.stopOpts
}
